"""
本例用于定义一个inception,是GoogLeNet中的一个多重复单元

目的:
inception的目的是卷积时自动卷积核(kernel,就是每种都试一试最后相加),
每个分支都使用不同的卷积核进行卷积,越有效的核,会在最终结果时具有更高的权重

实现方式:
inception具有多个分支,计算时将多个分支的结果,按照channel方向进行拼接
为了最终可以顺利拼接,inception的分支不会改变输入的width/height
"""
import torch
import torch.nn as nn


class InceptionSimple(nn.Module):
    # 定义inputChannelSize
    def __init__(self, inputChannelSize, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # 分之一
        # 经历Average Pooling池化层后,再经过1*1卷积得到输出
        # 定义Average Pooling池化层
        self.pool = nn.AvgPool2d(3, stride=1, padding=1)
        # 定义卷积层
        self.conv1 = nn.Conv2d(inputChannelSize, 24, 1)

        # 分之二
        # 只有1个卷积层,直接定义卷积层
        self.conv2 = nn.Conv2d(inputChannelSize, 16, 1)

        # 分之三
        # 先经过卷积层1,再经过卷积层2
        # 卷积层1和分支二的卷积层一样
        # 定义卷积层2
        self.conv3 = nn.Conv2d(16, 24, 5, padding=2)

        # 分之四
        # 先经过卷积层1->经过卷积层2->经过卷积层3
        # 卷积层1和分支二的卷积层一样
        # 定义卷积层2
        self.conv4 = nn.Conv2d(16, 24, 3, padding=1)
        # 定义卷积层3
        self.conv5 = nn.Conv2d(24, 24, 3, padding=1)

    def forward(self, x):
        # 分支一
        v1 = self.conv1(self.pool(x))
        # 分支二
        v2 = self.conv2(x)
        # 分支三
        v3 = self.conv3(self.conv2(x))
        # 分支四
        v4 = self.conv5(self.conv4(self.conv2(x)))

        # 维度分布为: patchSize,channel,h,w
        r = torch.cat((v1, v2, v3, v4), dim=1)
        return r


if __name__ == '__main__':
    a = torch.randn(1, 1, 28, 28)
    inception = InceptionSimple(1)
    out = inception(a)
    print(out.shape)
